import numpy as np

class EDA:
    '''
    EDA算法
    '''
    def __init__(self,number,size,dim):
        '''

        :param number: 迭代次数
        :param size: 粒子数
        :param dim: 决策维度
        '''
        self.number = number  #跌代次数
        self.dim = dim   #决策维度
        self.size = size  #粒子数
        self.fitNess = np.zeros(size) #适应值

        self.polulation = np.random.randint(0,2,self.size*self.dim).reshape(self.size,self.dim)
        self.port = 0.5 #择优占比
        # self.a = 0.3 #学习效率

        self.p = np.zeros(dim) #分布概率
        # self.p = np.ones(dim)*0.5

        self.bestPolu = np.zeros(int(self.size*self.port*self.dim)).reshape(int(self.size*self.port),self.dim)
        self.newPolu = np.zeros(int(self.size*(1-self.port))*self.dim).reshape(int(self.size*(1-self.port)),self.dim)
        #可以先使用p初始化第一代，而不是随机一代

        # #初始化种群
        # for i in range(self.size):
        #     p_temp = np.random.random(self.dim)
        #     for j in range(len(self.p)):
        #         if p_temp[j] <= self.p[j]:
        #             p_temp[j] = 1
        #         else:
        #             p_temp[j] = 0
        #     self.polulation[i] = p_temp

        #初始化fitNess
        # self.fitNess = self.getRsult(self.polulation)

        # #初始化p
        # self.bestPolu = self.select(self.fitNess)
        # for i in range(self.dim):
        #     self.p[i] = sum(self.bestPolu.T[i])/self.bestPolu.shape[0]

    def getRsult(self,polulation):
        #解码操作
        self.fitNess = np.zeros(self.size) #适应值
        for i in range(polulation.shape[1]):
            self.fitNess += (polulation.T[i, :] * (2 ** (polulation.shape[1]-1-i)))
        return self.fitNess

    def select(self,fitNess):
        '''
        根据fitness挑选优秀解
        :param fitNess:
        :return:
        '''
        tempIndex = np.argsort(fitNess)[-(int(len(fitNess)*self.port)):]
        for i in range(len(tempIndex)):
            self.bestPolu[i] = self.polulation[tempIndex[i]]
        return self.bestPolu

    def update(self):
        '''迭代更新'''
        for i in range(self.number):
            #计算适应值
            self.fitNess = self.getRsult(self.polulation)

            #选取优秀个体
            self.bestPolu = self.select(self.fitNess)
            #更新p
            for i in range(self.dim):
                self.p[i] = sum(self.bestPolu.T[i]) / self.bestPolu.shape[0]

                #使用学习效率来优化生成个体
                # self.p[i] = (1-self.a)*self[i]+self.a*sum(self.bestPolu.T[i]) / self.bestPolu.shape[0]

            #生成个体
            for i in range(self.size-self.bestPolu.shape[0]):
                p_temp = np.random.random(self.dim)
                for j in range(len(p_temp)):
                    if p_temp[j]<=self.p[j]:
                        p_temp[j] = 1
                    else:
                        p_temp[j] = 0
                self.newPolu[i] = p_temp

            #更新polulation
            self.polulation = np.vstack((self.bestPolu,self.newPolu))
            print(self.polulation)
            print()


if __name__ =="__main__":
    eda = EDA(20,30,10)
    eda.update()

